from torch import nn


class ConvBlock3d(nn.Module):
    def __init__(self, input_channel, output_channel, kernel_size=3, strides=1, padding=1, activation="relu", BN=True,
                 dropout=0.3, if_resnet=True, if_pooling=True, pooling_size=2, pooling_stride=2):
        super(ConvBlock3d, self).__init__()
        self.if_resnet = if_resnet
        self.all_operation = nn.Sequential()
        self.all_operation.add_module("conv1",
                                      nn.Conv3d(input_channel, output_channel, kernel_size, strides, padding))
        if activation.lower() == "relu":
            self.all_operation.add_module("active1", nn.ReLU())
        if BN:
            self.all_operation.add_module("BN1", nn.BatchNorm3d(output_channel))
        if dropout > 0:
            self.all_operation.add_module("dropout", nn.Dropout3d(p=dropout))
        self.all_operation.add_module("conv2",
                                      nn.Conv3d(output_channel, output_channel, kernel_size, strides, padding))
        if BN:
            self.all_operation.add_module("BN2", nn.BatchNorm3d(output_channel))
        self.if_pooling = if_pooling
        self.maxpooling = nn.MaxPool3d(kernel_size=pooling_size, stride=pooling_stride)


    def forward(self, x):
        if not self.if_resnet:
            x = self.all_operation(x)
            if self.if_pooling:
                x = self.maxpooling(x)
            return x
        else:
            residual = self.all_operation(x)
            if self.if_pooling:
                x = self.maxpooling(residual)
            return residual, x


class UpConvBlock3d(nn.Module):
    def __init__(self):
        super(UpConvBlock3d).__init__()
        self.upSampling = nn.Upsample()

    def forward(self, x):
        x = self.upSampling(x)
